package org.example.tools;

import dev.langchain4j.agent.tool.*;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.service.tool.DefaultToolExecutor;
import org.example.utils.ModelUtils;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

/**
 * @Author codeYang
 * @Date 2025/3/16 20:10
 * @Version 1.0
 */
public class ChatTools2 {

    static class WeatherUtils {

        @Tool("获取某个城市的天气")
        public String getWeather(@P("指定的城市") String city) {
            return "今天" + city + "天气不错";
        }

        @Tool("获取今天日期")
        public static String getData() {
            LocalDateTime now = LocalDateTime.now();
            DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
            return now.format(formatter);
        }
    }

    public static void main(String[] args) {

        ChatLanguageModel chatLanguageModel = ModelUtils.getZhipuAiModel();

        // 设置扩展工具类
        List<ToolSpecification> toolSpecifications = ToolSpecifications.toolSpecificationsFrom(WeatherUtils.class);

        // 创建消息队列
        List<ChatMessage> chatMessage = new ArrayList<>();

        UserMessage userMessage = UserMessage.from("今天日期是?");
        chatMessage.add(userMessage);

        // 构建工具与会话调用
        AiMessage aiMessage = chatLanguageModel.generate(chatMessage, toolSpecifications).content();
        // 本地执行方法
        if (aiMessage.hasToolExecutionRequests()) {
            for (ToolExecutionRequest toolExecutionRequest : aiMessage.toolExecutionRequests()) {
                String id = toolExecutionRequest.id();
                String name = toolExecutionRequest.name();
                String arguments = toolExecutionRequest.arguments();
                System.out.println("id:" + id + ",name:" + name + ",arg:" + arguments);
            }
        }
        // 回复结果
        chatMessage.add(aiMessage);


        WeatherUtils weatherUtils = new WeatherUtils();
        // 将工具调用方法与聊天消息一起传输个AI模型
        for (ToolExecutionRequest toolExecutionRequest : aiMessage.toolExecutionRequests()) {

            DefaultToolExecutor toolExecutor = new DefaultToolExecutor(weatherUtils, toolExecutionRequest);
            String result = toolExecutor.execute(toolExecutionRequest, UUID.randomUUID().toString().replace("-", ""));
            System.out.println("方法执行结果:" + result);
            // 传入id
            ToolExecutionResultMessage resultMessage = ToolExecutionResultMessage.from(toolExecutionRequest, result);

            chatMessage.add(resultMessage);
        }

        String text = chatLanguageModel.generate(chatMessage).content().text();
        System.out.println(text);

    }


}
